package middleware

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

type RateLimitMiddleware struct {
	FillInterval time.Duration
	Cap          int64
	QuanTum      int64
}

func NewRateLimitMiddleware(fillInterval time.Duration, cap, quantum int64) *RateLimitMiddleware {
	return &RateLimitMiddleware{
		FillInterval: fillInterval,
		Cap:          cap,
		QuanTum:      quantum,
	}
}

func (m *RateLimitMiddleware) Handle() gin.HandlerFunc {
	bucket := ratelimit.NewBucketWithQuantum(m.FillInterval, m.Cap, m.QuanTum)
	return func(c *gin.Context) {
		if bucket.TakeAvailable(1) < 1 {
			c.String(http.StatusForbidden, "Too many request")
			c.Abort()
			return
		}
		c.Next()
	}
}
